from mptt.utils import tree_item_iterator
from shop.models import Category, Product

def category_tree_for_node(current_category):
    root_categories = Category.tree.root_nodes()
    current_category_root = current_category.get_root()
    category_tree = []
    for category in root_categories:
        if current_category_root.id == category.id:
            descendants = category._tree_manager.add_related_count(category.get_descendants(True), Product, 'category', 'product_count', True)
            for descendant in descendants:
                if descendant.level <= current_category.level + 1:
                    category_tree.append(descendant)
        else:
            category_tree.append(category._tree_manager.add_related_count(category.get_descendants(True), Product, 'category', 'product_count', True)[0])
    return tree_item_iterator(category_tree)


